import argparse
import torch
import cv2
import numpy as np
from numpy import random
from models.experimental import attempt_load
from utils.general import check_img_size, check_requirements, check_imshow, non_max_suppression, scale_coords
from utils.torch_utils import select_device
from utils.datasets import LoadStreams, LoadImages
from utils.plots import plot_one_box

def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleup=True, stride=32):
    # 对输入图像进行前处理（做灰条填充）+变为tensor能认的4维
    # Resize and pad image while meeting stride-multiple constraints 调整大小和垫图像，同时满足跨步多约束
    shape = im.shape[:2]  # current shape [height, width]
    if isinstance(new_shape, int):
        new_shape = (new_shape, new_shape)

    # Scale ratio (new / old) 尺度比 (new / old)
    r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
    if not scaleup:  # only scale down, do not scale up (for better val mAP) # 只缩小，不扩大(为了更好的val mAP)
        r = min(r, 1.0)

    # Compute padding 计算填充
    new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
    dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding

    if auto:  # minimum rectangle 最小矩形区域
        dw, dh = np.mod(dw, stride), np.mod(dh, stride)  # wh padding

    dw /= 2  # divide padding into 2 sides
    dh /= 2

    if shape[::-1] != new_unpad:  # resize
        im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
    im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add border
    return im, r, (dw, dh)

def preprocess(img):
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # bgr->rgb
    image, ratio, dwdh = letterbox(img.copy(), auto=False)
    image = image.transpose((2, 0, 1))
    image = np.expand_dims(image, 0)
    image = np.ascontiguousarray(image)
    im = image.astype(np.float32)
    im /= 255  # (1, 3, 640, 640)
    return im, ratio, dwdh

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', nargs='+', type=str, default='/home/zhicheng.luo/Weights/chuwei_best_20240805.pt', help='model.pt path(s)')
    parser.add_argument('--source', type=str, default='/home/zhicheng.luo/data/chuwei/toilet_44.jpg', help='source')  # file/folder, 0 for webcam
    parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
    parser.add_argument('--conf-thres', type=float, default=0.35, help='object confidence threshold')
    parser.add_argument('--iou-thres', type=float, default=0.65, help='IOU threshold for NMS')
    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--view-img', action='store_true', default=True, help='display results')
    parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')
    parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
    parser.add_argument('--augment', action='store_true', help='augmented inference')
    parser.add_argument('--save-libtorch', type=bool, default=False,
                        help='save libtorch model for triton torch backend infer')#False#True
    opt = parser.parse_args()

    device = select_device(opt.device)
    model = attempt_load(opt.weights, map_location=device)  # load FP32 model
    model.eval()
    stride = int(model.stride.max())  # model stride
    imgsz = check_img_size(opt.img_size, s=stride)  # check img_size

    dataset = LoadImages(opt.source, img_size=imgsz, stride=stride)

    # Get names and colors
    names = model.module.names if hasattr(model, 'module') else model.names
    colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]

    for path, img, im0s, vid_cap in dataset:
        img = torch.from_numpy(img).to(device).float()
        img /= 255.0  # 0 - 255 to 0.0 - 1.0
        if img.ndimension() == 3:
            img = img.unsqueeze(0)

        if opt.save_libtorch:#如果是保存libtorch后端，那么输入图片的尺寸最好是固定的
            tmp = cv2.imread(opt.source)
            im, ratio, dwdh = preprocess(tmp)
            img = torch.from_numpy(im).to(device)

        with torch.no_grad():  # Calculating gradients would cause a GPU memory leak
            pred = model(img, augment=opt.augment)
            if isinstance(pred, tuple):
                pred = pred[0]
            pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
            pred_np = pred[0].cpu().numpy()

            # save for triton, 但是需要将models/yolo.py中Detect或IDetect类的成员变量concat改成True，因为triton不接受list[tensor]这样的输出
            if opt.save_libtorch:
                traced = torch.jit.trace(model, img)
                torch.jit.save(traced, opt.weights + "_model.pt")

            # Process detections
            for i, det in enumerate(pred):  # detections per image
                p, s, im0, frame = path, '', im0s, getattr(dataset, 'frame', 0)
                if len(det):
                    # Rescale boxes from img_size to im0 size
                    det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
                    for *xyxy, conf, cls in reversed(det):
                        if opt.view_img:  # Add bbox to image
                            label = f'{names[int(cls)]} {conf:.2f}'
                            plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=1)
                if opt.view_img:
                    # cv2.imshow(str(p), im0)
                    # cv2.waitKey(0)  # 1 millisecond
                    cv2.imwrite('/home/zhicheng.luo/data/chuwei/test/toilet_44_res.jpg', im0)
            print('over')